CUDA: support ncols > 1024 in bitonic argsort/top-k The original bito… by Geramy · Pull Request #26493 · ggml-org/llama.cpp

Geramy · 2026-08-03T05:43:35Z

Overview

This PR changes the argsort / top-k max cols to the maximum of the gpu available LDS by changing the kernel to use the maximum amount of X elements per thread. This was brought on due to the following issue:
#26399 (comment)

This mainly targets HIP devices due to there being no ggml_cuda_op_top_k implementation on hip.

Additional information

Please note drift is not yet calculated in the below table.

┌──────┬───────────┬──────────┬────────┐
│ n_kv │ TG before │ TG after │   %    │
├──────┼───────────┼──────────┼────────┤
│ 1004 │ 15.64     │ 15.13    │ −3.3%  │
├──────┼───────────┼──────────┼────────┤
│ 2004 │ 16.48     │ 16.29    │ −1.2%  │
├──────┼───────────┼──────────┼────────┤
│ 3004 │ 16.41     │ 16.24    │ −1.0%  │
├──────┼───────────┼──────────┼────────┤
│ 4004 │ 16.36     │ 16.18    │ −1.1%  │
├──────┼───────────┼──────────┼────────┤
│ 4504 │ 15.24     │ 16.05    │ +5.3%  │
├──────┼───────────┼──────────┼────────┤
│ 5004 │ 15.20     │ 16.08    │ +5.8%  │
├──────┼───────────┼──────────┼────────┤
│ 5504 │ 14.28     │ 16.02    │ +12.2% │
└──────┴───────────┴──────────┴────────┘

Requirements

2* AMD GPU
VRAM > 83GB
Deepseek-V4-Flash

I have read and agree with the contributing guidelines Yes
AI usage disclosure: YES, to run benchmarks.

…nic sort used when CUB is unavailable IE HIP, and some versions of older CUDA capped ncols to 1024 threads/blocks due to hardware LDS size limit and instead sent the request the the CPU. I have updated the kernel to now do multiple elements per thread ncols_pad/nthreads.

Geramy · 2026-08-03T05:46:46Z

@pwilkin

Geramy · 2026-08-03T16:30:18Z

This would close issue #26399 (comment)

remeh · 2026-08-04T08:10:05Z

This PR makes DeepSeek v4-Flash working fine on Strix Halo: there was a sudden drop of tg around 4k context (20->5 tok/s), which is not happening anymore with this PR merged. Thank you for work on this @Geramy 🙇

fairydreaming · 2026-08-04T12:55:40Z

I tested this on NVIDIA by adding:

#undef GGML_CUDA_USE_CUB
#undef CUB_TOP_K_AVAILABLE

in top-k.cu with my usual top-k torture set:

    for (auto nrows : {2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096 }) {
        for (auto cols : {128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65000, 200000}) {
            test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {cols, nrows, 1, 1}, 64));
         }
     }

added in test-backend-ops.cpp, however ./bin/test-backend-ops -o "TOP_K" crashed with:

/home/phm/projects/llama.cpp-hip-top-k-large-ncols/ggml/src/ggml-cuda/argsort.cu:242: GGML_ASSERT(shared_mem <= ggml_cuda_info().devices[ggml_cuda_get_device()].smpb) failed

Added printing shared mem usage:

$ ./bin/test-backend-ops -o "TOP_K"
ggml_cuda_init: found 1 CUDA devices (Total VRAM: 97247 MiB):
  Device 0: NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition, compute capability 12.0, VMM: yes, VRAM: 97247 MiB
Testing 2 devices

Backend 1/2: CUDA0
  Device description: NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition
  Device memory: 97247 MB (96640 MB free)

shared mem size: 512
  TOP_K(type=f32,ne=[128,2,1,1],k=64,ties=0): OK
shared mem size: 1024
  TOP_K(type=f32,ne=[256,2,1,1],k=64,ties=0): OK
shared mem size: 2048
ggml_backend_cuda_graph_compute: CUDA graph warmup complete
  TOP_K(type=f32,ne=[512,2,1,1],k=64,ties=0): OK
shared mem size: 4096
  TOP_K(type=f32,ne=[1024,2,1,1],k=64,ties=0): OK
shared mem size: 8192
  TOP_K(type=f32,ne=[2048,2,1,1],k=64,ties=0): OK
shared mem size: 16384
  TOP_K(type=f32,ne=[4096,2,1,1],k=64,ties=0): OK
shared mem size: 32768
  TOP_K(type=f32,ne=[8192,2,1,1],k=64,ties=0): OK
shared mem size: 65536

and it seems to grow with the row length. Not sure how it works on AMD, but doesn't seem right.

OK, before posting this I found your comment:

    // the shared memory budget is the limit on ncols_pad - keep this in sync with the check in
    // ggml_backend_cuda_device_supports_op() so that larger rows fall back instead of aborting

So @Geramy is aware of this. But not sure what's the point of this PR then. Slightly increasing supported row length before switching to CPU?

Ended up confused. What about kubakomu@42b262a, isn't this a more general AMD solution?

Geramy · 2026-08-04T14:55:13Z

I tested this on NVIDIA by adding:

#undef GGML_CUDA_USE_CUB
#undef CUB_TOP_K_AVAILABLE

in top-k.cu with my usual top-k torture set:

    for (auto nrows : {2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096 }) {
        for (auto cols : {128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65000, 200000}) {
            test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {cols, nrows, 1, 1}, 64));
         }
     }

added in test-backend-ops.cpp, however ./bin/test-backend-ops -o "TOP_K" crashed with:

/home/phm/projects/llama.cpp-hip-top-k-large-ncols/ggml/src/ggml-cuda/argsort.cu:242: GGML_ASSERT(shared_mem <= ggml_cuda_info().devices[ggml_cuda_get_device()].smpb) failed

Added printing shared mem usage:

$ ./bin/test-backend-ops -o "TOP_K"
ggml_cuda_init: found 1 CUDA devices (Total VRAM: 97247 MiB):
  Device 0: NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition, compute capability 12.0, VMM: yes, VRAM: 97247 MiB
Testing 2 devices

Backend 1/2: CUDA0
  Device description: NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition
  Device memory: 97247 MB (96640 MB free)

shared mem size: 512
  TOP_K(type=f32,ne=[128,2,1,1],k=64,ties=0): OK
shared mem size: 1024
  TOP_K(type=f32,ne=[256,2,1,1],k=64,ties=0): OK
shared mem size: 2048
ggml_backend_cuda_graph_compute: CUDA graph warmup complete
  TOP_K(type=f32,ne=[512,2,1,1],k=64,ties=0): OK
shared mem size: 4096
  TOP_K(type=f32,ne=[1024,2,1,1],k=64,ties=0): OK
shared mem size: 8192
  TOP_K(type=f32,ne=[2048,2,1,1],k=64,ties=0): OK
shared mem size: 16384
  TOP_K(type=f32,ne=[4096,2,1,1],k=64,ties=0): OK
shared mem size: 32768
  TOP_K(type=f32,ne=[8192,2,1,1],k=64,ties=0): OK
shared mem size: 65536

and it seems to grow with the row length. Not sure how it works on AMD, but doesn't seem right.

OK, before posting this I found your comment:

    // the shared memory budget is the limit on ncols_pad - keep this in sync with the check in
    // ggml_backend_cuda_device_supports_op() so that larger rows fall back instead of aborting

So @Geramy is aware of this. But not sure what's the point of this PR then. Slightly increasing supported row length before switching to CPU?

Ended up confused. What about kubakomu@42b262a, isn't this a more general AMD solution?

Technically it should also be working on nvidia, but this isn't a nvidia problem because llama.cpp doesn't use this method for nvidia. I can guard it with a HIP condition and put the original back for the nvidia fallback. I do appreciate you checking on nvidia, I believe nvidia is 48KB LDS I can also adjust that, but it maxes out around 12884 elements I believe or close to it.

IMbackK · 2026-08-04T15:00:44Z

probubly better to support the cub path on hip than trying to fit more into lds

Geramy · 2026-08-04T15:02:29Z

probubly better to support the cub path on hip than trying to fit more into lds

@IMbackK would you like me to close this PR and open one that implements CUB? I believe there are a few already but I can also do it too if it's preferred.

fairydreaming · 2026-08-04T15:29:51Z

@Geramy Since DeepSeek V4 has 1M max context length and ggml_top_k() is used in CSA attention that compresses this x4, to fully use the model on AMD cards without CPU offload we need ggml_top_k() to run on the GPU with row lengths at least up to 256k. GLM-5.2 and DeepSeek V3.2 use ggml_top_k() row lengths up to 1M as they use no cache compression.

So if we merged this PR as-is soon we would need another "CUDA: support ncols > 8192 in bitonic argsort/top-k" PR.

I searched for hipCUB PRs but only found #26388 that is closed unmerged. I don't know what are the hipCUB argsort row length limits, I guess it's something that still needs to be checked.

For me any solution is acceptable as long as it supports long row lengths, is reasonably fast and works correctly. hibCUB one seems easiest to implement.

Geramy · 2026-08-04T15:33:21Z

@fairydreaming @IMbackK

// top k elements per row
    // note: the resulting top k indices are in no particular order
    GGML_API struct ggml_tensor * ggml_top_k(
            struct ggml_context * ctx,
            struct ggml_tensor  * a,
            int

So technically instead of using DeviceSegmentedRadixSort::SortPairsDescending I could use rocprim::partition a radix-select instead, I would then be splitting top-k function and argsort function and we would have two implementations.

Neresco · 2026-08-04T16:35:44Z

I tried to use deepseek Flash v4 0731 since some days.
4x9060xt RDNA4 gfx1200 + 1x Strix halo gfx1151 over RPC.
This PR is the first that works for me without crashing.

PR 26388 had crashed too.

Thanks for this, from my non programmer eye's this is the right way because it works.

Neresco · 2026-08-04T16:58:12Z

Interestingly with this PR ROCM did crash to for me right now again, on the strix halo.
It happened while reusing the cache repeatedly.
Test without cache (-cram 0) is now ongoing.

32.04.881.699 I slot print_timing: id  1 | task 1059 | prompt eval time =  712561.81 ms / 62087 tokens (   11.48 ms per token,    87.13 tokens per second)
32.04.881.706 I slot print_timing: id  1 | task 1059 |        eval time =  760216.26 ms /   803 tokens (  946.72 ms per token,     1.06 tokens per second)
32.04.881.707 I slot print_timing: id  1 | task 1059 |       total time = 1472778.07 ms / 62890 tokens
32.04.881.708 I slot print_timing: id  1 | task 1059 |    graphs reused =       1822
32.04.883.648 I slot      release: id  1 | task 1059 | stop processing: n_tokens = 62889, truncated = 0
32.05.807.471 I slot get_availabl: id  1 | task -1 | selected slot by LCP similarity, f_sim_best = 0.990 (> 0.100 thold), f_keep = 0.987

32.05.807.549 I slot launch_slot_: id  1 | task 1931 | processing task, is_child = 0
32.11.721.152 I slot print_timing: id  1 | task 1931 | prompt processing, n_tokens =    384, progress = 1.00, t =   4.96 s / 77.46 tokens per second
32.16.865.606 I slot print_timing: id  1 | task 1931 | prompt processing, n_tokens =    653, progress = 1.00, t =  10.10 s / 64.64 tokens per second

32.26.199.879 I slot print_timing: id  1 | task 1931 | prompt eval time =   13743.31 ms /   657 tokens (   20.92 ms per token,    47.81 tokens per second)
32.26.199.883 I slot print_timing: id  1 | task 1931 |        eval time =    5693.01 ms /    42 tokens (  135.55 ms per token,     7.38 tokens per second)
32.26.199.884 I slot print_timing: id  1 | task 1931 |       total time =   19436.32 ms /   699 tokens
32.26.199.885 I slot print_timing: id  1 | task 1931 |    graphs reused =       1862
32.26.201.609 I slot      release: id  1 | task 1931 | stop processing: n_tokens = 62781, truncated = 0
32.26.786.868 I slot get_availabl: id  1 | task -1 | selected slot by LCP similarity, f_sim_best = 0.982 (> 0.100 thold), f_keep = 0.999

32.26.786.952 I slot launch_slot_: id  1 | task 1977 | processing task, is_child = 0
32.37.606.242 I slot print_timing: id  1 | task 1977 | prompt processing, n_tokens =    628, progress = 0.99, t =  10.68 s / 58.82 tokens per second
32.42.740.684 I slot print_timing: id  1 | task 1977 | prompt processing, n_tokens =   1140, progress = 1.00, t =  15.81 s / 72.10 tokens per second
33.01.982.917 I slot print_timing: id  1 | task 1977 | n_decoded =    100, tg =   7.37 t/s, tg_3s =   7.37 t/s
33.05.065.022 I slot print_timing: id  1 | task 1977 | n_decoded =    121, tg =   7.27 t/s, tg_3s =   6.81 t/s
33.07.363.597 I slot print_timing: id  1 | task 1977 | prompt eval time =   21492.30 ms /  1144 tokens (   18.79 ms per token,    53.23 tokens per second)
33.07.363.601 I slot print_timing: id  1 | task 1977 |        eval time =   18942.00 ms /   136 tokens (  139.28 ms per token,     7.18 tokens per second)
33.07.363.602 I slot print_timing: id  1 | task 1977 |       total time =   40434.30 ms /  1280 tokens
33.07.363.603 I slot print_timing: id  1 | task 1977 |    graphs reused =       1994
33.07.365.659 I slot      release: id  1 | task 1977 | stop processing: n_tokens = 64015, truncated = 0
33.07.946.204 I slot get_availabl: id  1 | task -1 | selected slot by LCP similarity, f_sim_best = 0.995 (> 0.100 thold), f_keep = 0.998
33.07.946.284 I slot launch_slot_: id  1 | task 2117 | processing task, is_child = 0
33.12.601.022 I slot print_timing: id  1 | task 2117 | prompt processing, n_tokens =    306, progress = 1.00, t =   4.53 s / 67.49 tokens per second

33.29.433.255 I slot print_timing: id  1 | task 2117 | n_decoded =    100, tg =   7.04 t/s, tg_3s =   7.04 t/s
33.32.545.620 I slot print_timing: id  1 | task 2117 | n_decoded =    123, tg =   7.11 t/s, tg_3s =   7.39 t/s
33.34.301.464 I slot print_timing: id  1 | task 2117 | prompt eval time =    7166.86 ms /   310 tokens (   23.12 ms per token,    43.25 tokens per second)
33.34.301.468 I slot print_timing: id  1 | task 2117 |        eval time =   19067.53 ms /   136 tokens (  140.20 ms per token,     7.13 tokens per second)
33.34.301.469 I slot print_timing: id  1 | task 2117 |       total time =   26234.39 ms /   446 tokens
33.34.301.470 I slot print_timing: id  1 | task 2117 |    graphs reused =       2125
33.34.303.283 I slot      release: id  1 | task 2117 | stop processing: n_tokens = 64321, truncated = 0
33.34.906.294 I slot get_availabl: id  1 | task -1 | selected slot by LCP similarity, f_sim_best = 0.958 (> 0.100 thold), f_keep = 0.998
33.34.906.368 I slot launch_slot_: id  1 | task 2256 | processing task, is_child = 0
33.46.753.228 I slot print_timing: id  1 | task 2256 | prompt processing, n_tokens =   1112, progress = 0.97, t =  11.73 s / 94.83 tokens per second
33.57.140.863 E recv failed (bytes_recv=0, size_to_recv=8)

Geramy · 2026-08-04T17:03:53Z

@Neresco are you building off of main?

Neresco · 2026-08-04T17:08:38Z

@Neresco are you building off of main?

No i directly downloaded your fork and compiled it.

Neresco · 2026-08-04T17:18:49Z

Ok it crashes now at ~65k instead of 4k it looks like.

 =  64270, progress = 0.96, t = 708.59 s / 90.70 tokens per second
18.30.362.339 I slot print_timing: id  3 | task 0 | prompt processing, n_tokens =  65294, progress = 0.97, t = 723.73 s / 90.22 tokens per second
18.42.093.273 E /home/lunarbuntu/Downloads/llama.cpp-Geramy-hip-top-k-large-ncols/ggml/src/ggml-rpc/ggml-rpc.cpp:519: Remote RPC server crashed or returned malformed response
recv failed (bytes_recv=0, size_to_recv=8)

Geramy · 2026-08-04T17:20:45Z

@Neresco cool! If you would like to test my other PR which is a draft right now, you may.
#26592

This other PR uses the CUB path instead on HIP and is the fix that @IMbackK actually wants.
I am still running it through final tests on my end.

Neresco · 2026-08-04T17:23:51Z

@Neresco cool! If you would like to test my other PR which is a draft right now, you may. #26592

This other PR uses the CUB path instead on HIP and is the fix that @IMbackK actually wants. I am still running it through final tests on my end.

Sure i am compiling and test. It will take ~45 minutes for the first result at ~65k threshold including compiling.

Neresco · 2026-08-04T18:14:43Z

@Neresco cool! If you would like to test my other PR which is a draft right now, you may. #26592

This other PR uses the CUB path instead on HIP and is the fix that @IMbackK actually wants. I am still running it through final tests on my end.

Two times now my Window Manager crashed around 63k.
By the second time i observed closely.
The Vram usage creeps up mb by mb until it crashed my window manager.
No llama.cpp crash like with the other PR.
I have now used some -ncmoe and test again with more headroom.

Geramy · 2026-08-04T18:27:29Z

@Neresco cool! If you would like to test my other PR which is a draft right now, you may. #26592
This other PR uses the CUB path instead on HIP and is the fix that @IMbackK actually wants. I am still running it through final tests on my end.

Two times now my Window Manager crashed around 63k. By the second time i observed closely. The Vram usage creeps up mb by mb until it crashed my window manager. No llama.cpp crash like with the other PR. I have now used some -ncmoe and test again with more headroom.

we should move off of this PR and move into #26592 I will be closing this PR once the other one goes from DRAFT to not.

Geramy · 2026-08-04T19:19:02Z

PR #26592 moved out of draft closing this PR.

Geramy requested a review from a team as a code owner August 3, 2026 05:43

github-actions Bot added ggml changes relating to the ggml tensor library for machine learning CUDA Related to the CUDA backend labels Aug 3, 2026

Geramy mentioned this pull request Aug 3, 2026

Eval bug: GGML_OP_TOP_K falls back to CPU on HIP/ROCm above ~3–4K context — 6.4× token-generation loss on DeepSeek-V4-Flash #26399

Open

IMbackK self-assigned this Aug 3, 2026

am17an requested a review from fairydreaming August 4, 2026 05:05

Geramy closed this Aug 4, 2026

remeh mentioned this pull request Aug 4, 2026

CUDA: enable the CUB path on HIP via hipCUB #26592

Open

Uh oh!

Conversation

Geramy commented Aug 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Overview

Additional information

Requirements

Uh oh!

Geramy commented Aug 3, 2026

Uh oh!

Geramy commented Aug 3, 2026

Uh oh!

remeh commented Aug 4, 2026

Uh oh!

fairydreaming commented Aug 4, 2026

Uh oh!

Geramy commented Aug 4, 2026

Uh oh!

IMbackK commented Aug 4, 2026

Uh oh!

Geramy commented Aug 4, 2026

Uh oh!

fairydreaming commented Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Geramy commented Aug 4, 2026

Uh oh!

Neresco commented Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Neresco commented Aug 4, 2026

Uh oh!

Geramy commented Aug 4, 2026

Uh oh!

Neresco commented Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Neresco commented Aug 4, 2026

Uh oh!

Geramy commented Aug 4, 2026

Uh oh!

Neresco commented Aug 4, 2026

Uh oh!

Neresco commented Aug 4, 2026

Uh oh!

Geramy commented Aug 4, 2026

Uh oh!

Geramy commented Aug 4, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Geramy commented Aug 3, 2026 •

edited

Loading

fairydreaming commented Aug 4, 2026 •

edited

Loading

Neresco commented Aug 4, 2026 •

edited

Loading

Neresco commented Aug 4, 2026 •

edited

Loading